import torch
import torch.nn as nn
import pandas as pd
from torch_geometric.nn import GCNConv
from torch_geometric.data import Data


class LightGCN(nn.Module):
    def __init__(self, num_users, num_items, embedding_dim, num_layers=2):
        super(LightGCN, self).__init__()
        self.num_users = num_users
        self.num_items = num_items
        self.embedding_dim = embedding_dim
        self.num_layers = num_layers

        # 初始化用户和物品的嵌入
        self.user_embedding = nn.Embedding(num_users, embedding_dim)
        self.item_embedding = nn.Embedding(num_items, embedding_dim)

        # GCN层
        self.convs = nn.ModuleList([GCNConv(embedding_dim, embedding_dim) for _ in range(num_layers)])

        # 初始化嵌入参数
        nn.init.xavier_uniform_(self.user_embedding.weight)
        nn.init.xavier_uniform_(self.item_embedding.weight)

    def forward(self, edge_index):
        # 初始嵌入
        x = torch.cat([self.user_embedding.weight, self.item_embedding.weight], dim=0)

        # 多轮GCN聚合
        embeddings = [x]
        for conv in self.convs:
            x = conv(x, edge_index)
            embeddings.append(x)

        # 取各层嵌入的平均值
        x = torch.mean(torch.stack(embeddings), dim=0)

        # 分离用户和物品嵌入
        user_emb = x[:self.num_users]
        item_emb = x[self.num_users:]
        return user_emb, item_emb


# 1. 从处理后的数据读取并构建交互图
# 读取处理后的评分数据
processed_ratings = pd.read_csv('processed_ratings.csv')

# 2. 数据预处理：构建用户-物品交互边
# 映射用户ID和物品ID到连续索引（LightGCN需要连续整数索引）
user_ids = processed_ratings['userId'].unique()
item_ids = processed_ratings['movieId'].unique()

user_id_map = {uid: i for i, uid in enumerate(user_ids)}
item_id_map = {mid: i for i, mid in enumerate(item_ids)}

num_users = len(user_ids)
num_items = len(item_ids)

# 构建边索引（用户索引 -> 物品索引，注意物品索引要偏移用户数量）
edges = []
for _, row in processed_ratings.iterrows():
    user_idx = user_id_map[row['userId']]
    item_idx = item_id_map[row['movieId']] + num_users  # 物品索引偏移用户数量，避免与用户索引冲突
    edges.append([user_idx, item_idx])

# 转换为PyTorch张量
edge_index = torch.tensor(edges, dtype=torch.long).t().contiguous()

# 3. 准备训练目标（使用评分作为目标值）
ratings = torch.tensor(processed_ratings['rating'].values, dtype=torch.float32)

# 4. 创建LightGCN模型
model = LightGCN(num_users=num_users, num_items=num_items, embedding_dim=64, num_layers=2)

# 5. 定义优化器和损失函数
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.MSELoss()

# 6. 训练模型
batch_size = 1024
num_epochs = 10

for epoch in range(num_epochs):
    model.train()
    optimizer.zero_grad()

    # 获取用户和物品嵌入
    user_emb, item_emb = model(edge_index)

    # 计算预测评分（用户嵌入与物品嵌入的内积）
    user_indices = torch.tensor([e[0] for e in edges], dtype=torch.long)
    item_indices = torch.tensor([e[1] - num_users for e in edges], dtype=torch.long)  # 还原物品索引

    pred_ratings = torch.sum(user_emb[user_indices] * item_emb[item_indices], dim=1)

    # 计算损失
    loss = criterion(pred_ratings, ratings)
    loss.backward()
    optimizer.step()

    print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {loss.item():.4f}")

# 保存模型
torch.save(model.state_dict(), 'lightgcn_recommender.pth')
print("模型已保存为 lightgcn_recommender.pth")
